/*
 * Copyright 2002-2019 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.core.io.buffer;

import io.netty.buffer.PooledByteBufAllocator;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.jetbrains.annotations.NotNull;
import org.junit.After;
import org.springframework.util.Assert;

import java.nio.ByteBuffer;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * Implementation of the {@code DataBufferFactory} interface that keep track of memory leaks.
 * Useful for unit tests that handle data buffers. Simply inherit from
 * {@link AbstractLeakCheckingTestCase} or call {@link #checkForLeaks()} in
 * a JUnit {@link After} method yourself, and any buffers have not been released will result in an
 * {@link AssertionError}.
 *
 * @author Arjen Poutsma
 * @see LeakAwareDataBufferFactory
 */
public class LeakAwareDataBufferFactory implements DataBufferFactory {

    private static final Log logger = LogFactory.getLog(LeakAwareDataBufferFactory.class);


    private final DataBufferFactory delegate;

    private final List<LeakAwareDataBuffer> created = new ArrayList<>();


    /**
     * Creates a new {@code LeakAwareDataBufferFactory} by wrapping a
     * {@link DefaultDataBufferFactory}.
     */
    public LeakAwareDataBufferFactory() {
        this(new NettyDataBufferFactory(PooledByteBufAllocator.DEFAULT));
    }

    /**
     * Creates a new {@code LeakAwareDataBufferFactory} by wrapping the given delegate.
     *
     * @param delegate the delegate buffer factory to wrap.
     */
    public LeakAwareDataBufferFactory(DataBufferFactory delegate) {
        Assert.notNull(delegate, "Delegate must not be null");
        this.delegate = delegate;
    }


    /**
     * Checks whether all of the data buffers allocated by this factory have also been released.
     * If not, then an {@link AssertionError} is thrown. Typically used from a JUnit {@link After}
     * method.
     */
    public void checkForLeaks() {
        Instant start = Instant.now();
        while (true) {
            if (this.created.stream().noneMatch(LeakAwareDataBuffer::isAllocated)) {
                return;
            }
            if (Instant.now().isBefore(start.plus(Duration.ofSeconds(5)))) {
                try {
                    Thread.sleep(50);
                }
                catch (InterruptedException ex) {
                    // ignore
                }
                continue;
            }
            List<AssertionError> errors = this.created.stream()
                    .filter(LeakAwareDataBuffer::isAllocated)
                    .map(LeakAwareDataBuffer::leakError)
                    .collect(Collectors.toList());

            errors.forEach(it -> logger.error("Leaked error: ", it));
            throw new AssertionError(errors.size() + " buffer leaks detected (see logs above)");
        }
    }

    @Override
    public DataBuffer allocateBuffer() {
        return createLeakAwareDataBuffer(this.delegate.allocateBuffer());
    }

    @Override
    public DataBuffer allocateBuffer(int initialCapacity) {
        return createLeakAwareDataBuffer(this.delegate.allocateBuffer(initialCapacity));
    }

    @NotNull
    private DataBuffer createLeakAwareDataBuffer(DataBuffer delegateBuffer) {
        LeakAwareDataBuffer dataBuffer = new LeakAwareDataBuffer(delegateBuffer, this);
        this.created.add(dataBuffer);
        return dataBuffer;
    }

    @Override
    public DataBuffer wrap(ByteBuffer byteBuffer) {
        return this.delegate.wrap(byteBuffer);
    }

    @Override
    public DataBuffer wrap(byte[] bytes) {
        return this.delegate.wrap(bytes);
    }

    @Override
    public DataBuffer join(List<? extends DataBuffer> dataBuffers) {
        // Remove LeakAwareDataBuffer wrapper so delegate can find native buffers
        dataBuffers = dataBuffers.stream()
                .map(o -> o instanceof LeakAwareDataBuffer ? ((LeakAwareDataBuffer) o).getDelegate() : o)
                .collect(Collectors.toList());
        return new LeakAwareDataBuffer(this.delegate.join(dataBuffers), this);
    }

}
